package com.example.gateway.filter;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;

public class JwtGenerator {

    private static final String SECRET_KEY = "your-256-bit-secret"; // 256-bit secret key

    public static String generateJwt(String subject, long expirationTimeMs) {
        // Step 1: Create the header
        Map<String, Object> header = new HashMap<>();
        header.put("alg", "HS256");
        header.put("typ", "JWT");

        // Step 2: Create the payload
        Map<String, Object> payload = new HashMap<>();
        payload.put("sub", subject);
        payload.put("iss", "example.com");
        payload.put("exp", new Date(System.currentTimeMillis() + expirationTimeMs).getTime() / 1000);

        // Step 3: Encode the header and payload to Base64Url
        String headerJson = mapToJson(header);
        String payloadJson = mapToJson(payload);
        String base64UrlHeader = base64UrlEncode(headerJson.getBytes(StandardCharsets.UTF_8));
        String base64UrlPayload = base64UrlEncode(payloadJson.getBytes(StandardCharsets.UTF_8));

        // Step 4: Create the signature
        String data = base64UrlHeader + "." + base64UrlPayload;
        String signature = hmacSha256(data, SECRET_KEY);
        String base64UrlSignature = base64UrlEncode(signature.getBytes(StandardCharsets.UTF_8));

        // Step 5: Combine the header, payload, and signature
        return base64UrlHeader + "." + base64UrlPayload + "." + base64UrlSignature;
    }

    private static String mapToJson(Map<String, Object> map) {
        StringBuilder json = new StringBuilder();
        json.append("{");
        for (Map.Entry<String, Object> entry : map.entrySet()) {
            json.append("\"").append(entry.getKey()).append("\":\"");
            json.append(entry.getValue()).append("\",");
        }
        if (json.length() > 1) {
            json.setLength(json.length() - 1); // Remove the last comma
        }
        json.append("}");
        return json.toString();
    }

    private static String base64UrlEncode(byte[] data) {
        return Base64.getUrlEncoder().withoutPadding().encodeToString(data);
    }

    private static String hmacSha256(String data, String key) {
        try {
            Mac sha256Hmac = Mac.getInstance("HmacSHA256");
            SecretKeySpec secretKey = new SecretKeySpec(key.getBytes(StandardCharsets.UTF_8), "HmacSHA256");
            sha256Hmac.init(secretKey);
            byte[] hash = sha256Hmac.doFinal(data.getBytes(StandardCharsets.UTF_8));
            return Base64.getEncoder().encodeToString(hash);
        } catch (Exception e) {
            throw new RuntimeException("Failed to calculate HMAC-SHA256", e);
        }
    }

    public static void main(String[] args) {
        String subject = "user123";
        long expirationTimeMs = 3600000; // 1 hour
        String jwt = generateJwt(subject, expirationTimeMs);
        System.out.println("Generated JWT: " + jwt);
    }
}

